package org.apache.hadoop.hive.ql.hooks;

import com.sql.blood.common.util.GraphDatabaseConnection;
import com.sql.blood.common.model.TabLine;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.ql.QueryPlan;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.List;

public class HiveLineage implements ExecuteWithHookContext {

    private final static Logger logger = LoggerFactory.getLogger(HiveLineage.class);

    public void run(HookContext hookContext) throws Exception {
        logger.info("execute hook ...");
        assert (hookContext.getHookType() == HookContext.HookType.POST_EXEC_HOOK);
        QueryPlan queryPlan = hookContext.getQueryPlan();
        HiveConf conf = hookContext.getConf();
        String queryStr = queryPlan.getQueryStr();
        HiveParse hiveParse = new HiveParse();
        hiveParse.setUrl(conf.getVar(HiveConf.ConfVars.METASTOREURIS));
        hiveParse.getLineageInfo(queryStr);
        List<TabLine> input = hiveParse.getInputTableList();
        List<TabLine> output = hiveParse.getOutputTableList();
        GraphDatabaseConnection conn = new GraphDatabaseConnection();
        for (TabLine tabLine:input){
            conn.insertVertex(tabLine);
        }
        for (TabLine tabLine:output){
            conn.insertVertex(tabLine);
        }
        for (TabLine x:input){
            for (TabLine y:output){
                conn.insertEdge(x,y);
            }
        }
        conn.close();
    }
}
